{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FastPitch Speaker Interpolation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this tutorial we will explore the following:\n",
    "- What is speaker interpolation in a multi-speaker Text-to-Speech model\n",
    "- Why we need speaker interpolation\n",
    "- Step-by-step demo of performing speaker interpolation using pretrained multi-speaker FastPitch model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## License\n",
    "\n",
    "> Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.\n",
    ">\n",
    "> Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "> you may not use this file except in compliance with the License.\n",
    "> You may obtain a copy of the License at\n",
    ">\n",
    ">     http://www.apache.org/licenses/LICENSE-2.0\n",
    ">\n",
    "> Unless required by applicable law or agreed to in writing, software\n",
    "> distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "> WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "> See the License for the specific language governing permissions and\n",
    "> limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### What is Speaker Interpolation in a Multi-speaker Text-to-speech Model\n",
    "\n",
    "Given a Multi-speaker Text-to-Speech (TTS) model that can generate speech in more than one speaker's voice, Speaker Interpolation is the process of synthesizing new synthetic **voices by combining two or more existing speaker voices** without the need for any further finetuning. Before we go in depths of the process let's understand why we need speaker interpolation?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Why we need speaker interpolation\n",
    "\n",
    "To release any TTS model, trained on datasets where the speakers do not want speech to be generated in their voice, speaker interpolation is important. This technique can be used to replace the original speaker's in the pretrained model with the interpolated synthetic speakers. This work would solve this problem.\n",
    "\n",
    "\n",
    "This technique can be used to generate large multi-speaker dataset for training Automatic Speech Recognition (ASR) models or voice conversion models. It enables users to create an infinite number of speaker voices. \n",
    "\n",
    "We can also use speaker interpolation for achieving varying multi-dialect or multi-emotional voices in text-to-speech (TTS) synthesis. They may\n",
    "be used for personalizing speech and generating speech with specific voice characteristics in speech synthesis systems.\n",
    "\n",
    "In case of self-supervised learning (SSL), speaker interpolation can be used for negative sampling. In addition to the above use cases, speaker interpolation can be used for data augmentation as well."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step-by-step demo of Speaker Interpolation using Pretrained multi-speaker FastPitch model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Method\n",
    "\n",
    "1. Train a multi-speaker spectrogram generator model (like FastPitch) using embedding lookup table. This means that during training an embedding (speaker embedding) will be learnt for each speaker.\n",
    "\n",
    "2. Extract speaker embeddings for any two speakers, you want to combine, from the model trained in step 1.\n",
    "\n",
    "3. Perform a weighted sum of the two speaker embeddings extracted in step 2. This will give the speaker embedding for the interpolated speaker.\n",
    "\n",
    "4. Use the speaker embedding obtained from step 3, to condition the multi-speaker model in step 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install NeMo library. If you are running locally (rather than on Google Colab), comment out the below lines\n",
    "# and instead follow the instructions at https://github.com/NVIDIA/NeMo#Installation\n",
    "BRANCH = 'main'\n",
    "!python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[all]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# If seaborn is not installed already run this cell\n",
    "!pip install seaborn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "You can either run this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\n",
    "Instructions for setting up Colab are as follows:\n",
    "1. Open a new Python 3 notebook.\n",
    "2. Import this notebook from GitHub (File -> Upload Notebook -> \"GITHUB\" tab -> copy/paste GitHub URL)\n",
    "3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select \"GPU\" for hardware accelerator)\n",
    "4. Run this cell to set up dependencies# .\n",
    "\"\"\"\n",
    "\n",
    "import json\n",
    "\n",
    "import torch\n",
    "import numpy as np\n",
    "import IPython.display as ipd\n",
    "from matplotlib.pyplot import imshow\n",
    "from matplotlib import pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "from nemo.collections.tts.models import FastPitchModel, HifiGanModel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Load Pretrained multi-speaker FastPitch model**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's view the available model checkpoints for FastPitch \n",
    "FastPitchModel.list_available_models()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We will choose the English multi-speaker model: tts_en_fastpitch_multispeaker\n",
    "spec_gen_model = FastPitchModel.from_pretrained(\"tts_en_fastpitch_multispeaker\")\n",
    "spec_gen_model.cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Visualization**\n",
    "\n",
    "To understand the process better, let's try to visualize how the embeddings get combined to create a new speaker embedding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Original Speaker IDs\n",
    "speaker_1 = 2\n",
    "speaker_2 = 10\n",
    "\n",
    "# Extracting speaker embeddings for speaker_1 and speaker_2\n",
    "speaker_emb_1 = spec_gen_model.fastpitch.speaker_emb(torch.tensor(speaker_1, dtype=torch.int32).cuda()).clone().detach()\n",
    "speaker_emb_2 = spec_gen_model.fastpitch.speaker_emb(torch.tensor(speaker_2, dtype=torch.int32).cuda()).clone().detach()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will now visualize the two speaker embeddings. The embedding dimensions are 384, so to keep the visualization clear we will only visualize the first 30 elements of each embedding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "plt.rcParams['figure.figsize'] = [20, 1]\n",
    "tmp_emb = np.vstack([speaker_emb_1.cpu()[:30], speaker_emb_2.cpu()[:30]])\n",
    "sns.heatmap(data=tmp_emb, cmap=\"coolwarm\", xticklabels=False, yticklabels=[\"speaker1\", \"speaker2\"])\n",
    "plt.title(\"Visualizing Speaker Embeddings\", fontsize =20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's combine the two speaker embeddings and generate the synthetic speaker embedding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Weights assigned to each speaker, in the process of generating the new speaker\n",
    "# You can play around with these weights\n",
    "weight_speaker_1 = 0.5\n",
    "weight_speaker_2 = 0.5\n",
    "\n",
    "interpolated_speaker_emb = weight_speaker_1 * speaker_emb_1 + weight_speaker_2 * speaker_emb_2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will now visualize the interpolated speaker along with the two speaker original embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "plt.rcParams['figure.figsize'] = [20, 2]\n",
    "tmp_emb = np.vstack([speaker_emb_1.cpu()[:30], speaker_emb_2.cpu()[:30], interpolated_speaker_emb.cpu()[:30]])\n",
    "sns.heatmap(data=tmp_emb, cmap=\"coolwarm\", xticklabels=False, yticklabels=[\"speaker1\", \"speaker2\", \"interpolated_speaker\"])\n",
    "plt.title(\"Visualizing Original and Interpolated Speaker Embeddings\", fontsize =20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Extraction and Combining of speaker embeddings**\n",
    "\n",
    "This entire process of doing speaker interpolation is performed by the method `interpolate_speaker` in the FastPitch class."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Generate Speech using speaker interpolation**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we will use this process of speaker interpolation to generate speech using original speaker voices and interpolated speaker voice. For this we will need a vocoder to convert generated spectrogram to waveform speech. So let's load HiFiGAN vocoder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vocoder = HifiGanModel.from_pretrained(\"tts_en_hifitts_hifigan_ft_fastpitch\")\n",
    "vocoder.cuda().eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we will define a function which will take as input the spectrogram generator model, vocoder model, text to be converted to speech and speaker id. This function will return spectrogram and audio of the generated speech."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def infer(spec_gen_model, vocoder_model, str_input, speaker=None):\n",
    "    \"\"\"\n",
    "    Synthesizes spectrogram and audio from a text string given a spectrogram synthesis and vocoder model.\n",
    "    \n",
    "    Args:\n",
    "        spec_gen_model: Spectrogram generator model (FastPitch in our case)\n",
    "        vocoder_model: Vocoder model (HiFiGAN in our case)\n",
    "        str_input: Text input for the synthesis\n",
    "        speaker: Speaker ID\n",
    "    \n",
    "    Returns:\n",
    "        spectrogram and waveform of the synthesized audio.\n",
    "    \"\"\"\n",
    "    spec_gen_model.eval()\n",
    "    with torch.no_grad():\n",
    "        parsed = spec_gen_model.parse(str_input)\n",
    "        if speaker is not None:\n",
    "            speaker = torch.tensor([speaker]).long().to(device=spec_gen_model.device)\n",
    "        spectrogram = spec_gen_model.generate_spectrogram(tokens=parsed, speaker=speaker)\n",
    "        audio = vocoder_model.convert_spectrogram_to_audio(spec=spectrogram)\n",
    "        \n",
    "    if spectrogram is not None:\n",
    "        if isinstance(spectrogram, torch.Tensor):\n",
    "            spectrogram = spectrogram.to('cpu').numpy()\n",
    "        if len(spectrogram.shape) == 3:\n",
    "            spectrogram = spectrogram[0]\n",
    "    else:\n",
    "        raise Exception(\"None value was generated for spectrogram\")\n",
    "    if isinstance(audio, torch.Tensor):\n",
    "        audio = audio.to('cpu').numpy()\n",
    "    return spectrogram, audio"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's generate some speech."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Choose any original speakers (speaker1 and speaker2)\n",
    "# Choose new_speaker ID such that it is not already assigned to an existing speaker\n",
    "speaker1, speaker2, new_speaker = 6097, 92, 1000\n",
    "\n",
    "# Choose weight1 corresponding to speaker1 and weight2 corresponding to speaker2\n",
    "# You can change this and play around with the weights\n",
    "weight1, weight2 = 0.4, 0.6\n",
    "\n",
    "# Perform Speaker Interpolation\n",
    "spec_gen_model.interpolate_speaker(speaker1, speaker2, weight1, weight2, new_speaker)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is the text you want to convert to speech\n",
    "text = \"Hello World, Speaker Interpolation is cool\"\n",
    "\n",
    "_, audio = infer(spec_gen_model, vocoder, text, speaker=speaker1)\n",
    "print(f\"audio generated for speaker {speaker1}\")\n",
    "ipd.display(ipd.Audio(audio, rate=44100))\n",
    "\n",
    "_, audio = infer(spec_gen_model, vocoder, text, speaker=speaker2)\n",
    "print(f\"audio generated for speaker {speaker2}\")\n",
    "ipd.display(ipd.Audio(audio, rate=44100))\n",
    "\n",
    "_, audio = infer(spec_gen_model, vocoder, text, speaker=new_speaker)\n",
    "print(f\"audio generated for Interpolated Speaker\")\n",
    "ipd.display(ipd.Audio(audio, rate=44100))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It was observed that for speaker interpolation to work well and stable, the pretrained multi-speaker FastPitch model needs to be robust. So, training multi-speaker FastPitch model on larger data with larger number of speakers help in getting better interpolated speakers."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
